
import cv2
import matplotlib.pyplot as plt
from ultralytics import YOLO
import numpy as np

# 读取图像，以灰度模式读取
color_image = cv2.imread('test.png')
model = YOLO('../ultralytics/assets/yolo11n-seg.pt')
# 使用 YOLOv8 进行目标检测
# 这种方式不支持灰度图像 需要使用三通道的彩色图像
results = model(color_image)
# 获取检测到的人的掩码
person_masks = []
for result in results:
    masks = result.masks
    boxes = result.boxes
    for i, box in enumerate(boxes):
        if int(box.cls[0]) == 0:  # 类别 0 表示人
            person_masks.append(masks.data[i].cpu().numpy())

# 对每个人的区域进行处理
for i, mask in enumerate(person_masks):
    # 调整掩码的大小以匹配原始图像
    mask = cv2.resize(mask.astype(np.uint8), (color_image.shape[1], color_image.shape[0]))

    # 提取人的区域
    person = cv2.bitwise_and(color_image, color_image, mask=mask)

    # 将人的区域转换为灰度图
    gray_person = cv2.cvtColor(person, cv2.COLOR_BGR2GRAY)

    # 高斯模糊，减少噪声
    blurred = cv2.GaussianBlur(gray_person, (5, 5), 0)

    # 直方图均衡化，增强对比度
    equalized = cv2.equalizeHist(blurred)

    # 边缘检测
    edges = cv2.Canny(equalized, 50, 150)

    # 获取人的轮廓
    contours, _ = cv2.findContours(edges.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if contours:  # 检查是否检测到轮廓
        # 绘制人物轮廓
        cv2.drawContours(color_image, contours, -1, (255, 0, 0), 2)
    else:
        print(f"未检测到 Person {i + 1} 的轮廓")

# 显示分割后的图像
plt.figure(figsize=(10, 10))
plt.imshow(cv2.cvtColor(color_image, cv2.COLOR_BGR2RGB))
plt.title('Person Segmentation with Contours')
plt.axis('off')
plt.show()